// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <stdlib.h>
#include <unistd.h>
#include <stdio.h>
#include <signal.h>
#include <ktest/ktest.h>
#include <turbo/base/compat.h>
#include <turbo/times/time.h>
#include <turbo/base/macros.h>
#include <turbo/utility/errno.h>
#include <kthread/internal/sys_futex.h>
#include <kthread/internal/kutex.h>
#include <kthread/internal/kthread.h>
#include <atomic>
#include <turbo/flags/flag.h>
#include <cinttypes>

TURBO_FLAG(int32_t, thread_num, 1, "#pairs of threads doing ping pong");
TURBO_FLAG(bool, loop, false, "run until ctrl-C is pressed");
TURBO_FLAG(bool, use_futex, false, "use futex instead of pipe");
TURBO_FLAG(bool, use_kutex, false, "use kutex instead of pipe");

namespace {

    void ALLOW_UNUSED (*ignore_sigpipe)(int) = signal(SIGPIPE, SIG_IGN);

    volatile bool stop = false;

    void quit_handler(int) {
        stop = true;
    }

    struct TURBO_CACHELINE_ALIGNED AlignedIntWrapper {
        int value;
    };

    struct TURBO_CACHELINE_ALIGNED PlayerArg {
        int read_fd;
        int write_fd;
        int *wait_addr;
        int *wake_addr;
        long counter;
        long wakeup;
    };

    void *pipe_player(void *void_arg) {
        PlayerArg *arg = static_cast<PlayerArg *>(void_arg);
        char dummy = '\0';
        while (1) {
            ssize_t nr = read(arg->read_fd, &dummy, 1);
            if (nr <= 0) {
                if (nr == 0) {
                    printf("[%"
                           PRIu64
                           "] EOF\n", pthread_numeric_id());
                    break;
                }
                if (errno != EINTR) {
                    printf("[%"
                           PRIu64
                           "] bad read, %m\n", pthread_numeric_id());
                    break;
                }
                continue;
            }
            if (1L != write(arg->write_fd, &dummy, 1)) {
                printf("[%"
                       PRIu64
                       "] bad write, %m\n", pthread_numeric_id());
                break;
            }
            ++arg->counter;
        }
        return nullptr;
    }

    static const int INITIAL_FUTEX_VALUE = 0;

    void *futex_player(void *void_arg) {
        PlayerArg *arg = static_cast<PlayerArg *>(void_arg);
        int counter = INITIAL_FUTEX_VALUE;
        while (!stop) {
            int rc = kthread::futex_wait_private(arg->wait_addr, counter, nullptr);
            ++counter;
            ++*arg->wake_addr;
            kthread::futex_wake_private(arg->wake_addr, 1);
            ++arg->counter;
            arg->wakeup += (rc == 0);
        }
        return nullptr;
    }

    void *kutex_player(void *void_arg) {
        PlayerArg *arg = static_cast<PlayerArg *>(void_arg);
        int counter = INITIAL_FUTEX_VALUE;
        while (!stop) {
            int rc = kthread::kutex_wait(arg->wait_addr, counter, nullptr);
            ++counter;
            ++*arg->wake_addr;
            kthread::kutex_wake(arg->wake_addr);
            ++arg->counter;
            arg->wakeup += (rc == 0);
        }
        return nullptr;
    }

    TEST(PingPongTest, ping_pong) {
        signal(SIGINT, quit_handler);
        stop = false;
        PlayerArg *args[turbo::get_flag(FLAGS_thread_num)];

        for (int i = 0; i < turbo::get_flag(FLAGS_thread_num); ++i) {
            int pipe1[2];
            int pipe2[2];
            if (!turbo::get_flag(FLAGS_use_futex) && !turbo::get_flag(FLAGS_use_kutex)) {
                ASSERT_EQ(0, pipe(pipe1));
                ASSERT_EQ(0, pipe(pipe2));
            }

            PlayerArg *arg1 = new PlayerArg;
            if (!turbo::get_flag(FLAGS_use_futex) && !turbo::get_flag(FLAGS_use_kutex)) {
                arg1->read_fd = pipe1[0];
                arg1->write_fd = pipe2[1];
            } else if (turbo::get_flag(FLAGS_use_futex)) {
                AlignedIntWrapper *w1 = new AlignedIntWrapper;
                w1->value = INITIAL_FUTEX_VALUE;
                AlignedIntWrapper *w2 = new AlignedIntWrapper;
                w2->value = INITIAL_FUTEX_VALUE;
                arg1->wait_addr = &w1->value;
                arg1->wake_addr = &w2->value;
            } else if (turbo::get_flag(FLAGS_use_kutex)) {
                arg1->wait_addr = kthread::kutex_create_checked<int>();
                *arg1->wait_addr = INITIAL_FUTEX_VALUE;
                arg1->wake_addr = kthread::kutex_create_checked<int>();
                *arg1->wake_addr = INITIAL_FUTEX_VALUE;
            } else {
                ASSERT_TRUE(false);
            }
            arg1->counter = 0;
            arg1->wakeup = 0;
            args[i] = arg1;

            PlayerArg *arg2 = new PlayerArg;
            if (!turbo::get_flag(FLAGS_use_futex) && !turbo::get_flag(FLAGS_use_kutex)) {
                arg2->read_fd = pipe2[0];
                arg2->write_fd = pipe1[1];
            } else {
                arg2->wait_addr = arg1->wake_addr;
                arg2->wake_addr = arg1->wait_addr;
            }
            arg2->counter = 0;
            arg2->wakeup = 0;

            pthread_t th1, th2;
            kthread_t bth1, bth2;
            if (!turbo::get_flag(FLAGS_use_futex) && !turbo::get_flag(FLAGS_use_kutex)) {
                ASSERT_EQ(0, pthread_create(&th1, nullptr, pipe_player, arg1));
                ASSERT_EQ(0, pthread_create(&th2, nullptr, pipe_player, arg2));
            } else if (turbo::get_flag(FLAGS_use_futex)) {
                ASSERT_EQ(0, pthread_create(&th1, nullptr, futex_player, arg1));
                ASSERT_EQ(0, pthread_create(&th2, nullptr, futex_player, arg2));
            } else if (turbo::get_flag(FLAGS_use_kutex)) {
                ASSERT_EQ(0, kthread_start_background(&bth1, nullptr, kutex_player, arg1));
                ASSERT_EQ(0, kthread_start_background(&bth2, nullptr, kutex_player, arg2));
            } else {
                ASSERT_TRUE(false);
            }

            if (!turbo::get_flag(FLAGS_use_futex) && !turbo::get_flag(FLAGS_use_kutex)) {
                // send the seed data.
                unsigned char seed = 255;
                ASSERT_EQ(1L, write(pipe1[1], &seed, 1));
            } else if (turbo::get_flag(FLAGS_use_futex)) {
                ++*arg1->wait_addr;
                kthread::futex_wake_private(arg1->wait_addr, 1);
            } else if (turbo::get_flag(FLAGS_use_kutex)) {
                ++*arg1->wait_addr;
                kthread::kutex_wake(arg1->wait_addr);
            } else {
                ASSERT_TRUE(false);
            }
        }

        long last_counter = 0;
        long last_wakeup = 0;
        while (!stop) {
            turbo::TimeCost tm;
            tm.reset();
            sleep(1);
            tm.stop();
            long cur_counter = 0;
            long cur_wakeup = 0;
            for (int i = 0; i < turbo::get_flag(FLAGS_thread_num); ++i) {
                cur_counter += args[i]->counter;
                cur_wakeup += args[i]->wakeup;
            }
            if (turbo::get_flag(FLAGS_use_futex) || turbo::get_flag(FLAGS_use_kutex)) {
                printf("pingpong-ed %"
                       PRId64
                       "/s, wakeup=%"
                       PRId64
                       "/s\n",
                       (cur_counter - last_counter) * 1000L / tm.m_elapsed(),
                       (cur_wakeup - last_wakeup) * 1000L / tm.m_elapsed());
            } else {
                printf("pingpong-ed %"
                       PRId64
                       "/s\n",
                       (cur_counter - last_counter) * 1000L / tm.m_elapsed());
            }
            last_counter = cur_counter;
            last_wakeup = cur_wakeup;
            if (!turbo::get_flag(FLAGS_loop)) {
                break;
            }
        }
        stop = true;
        // Program quits, Let resource leak.
    }
} // namespace
